import ollama
import gradio as gr

# 文档处理和检索
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma

# 嵌入生成
from langchain_ollama.embeddings import OllamaEmbeddings

import re

""" response = ollama.chat(
    model = "deepseek-r1:1.5b",
    messages = [
        {"role": "user", "content":"解释牛顿第二运动定律"},
    ],
)

print(response["message"]["content"]) """

def process_pdf(pdf_bytes):
    if pdf_bytes is None:
        return None, None, None
    loader = PyMuPDFLoader(pdf_bytes)
    data = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size = 500, chunk_overlap = 100)
    chunks = text_splitter.split_documents(data)
    embeddings = OllamaEmbeddings(model = 'deepseek-r1:1.5b')
    vectorstores = Chroma.from_documents(documents = chunks, embedding = embeddings, persist_directory="./chroma_db")
    retriever = vectorstores.as_retriever()

    return text_splitter, vectorstores, retriever

def combine_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

def ollama_llm(question, context):
    formatted_prompt = f"问题：{question} \n\n 环境：{context}"
    response = ollama.chat(
        model = "deepseek-r1:1.5b",
        
        messages = [{"role":"user", "content": formatted_prompt}]
    )
    response_content = response["message"]["content"]
    final_answer = re.sub(r'<think>.*?</think>',
    '',response_content, flags=re.DOTALL).strip()
    return final_answer

def rag_chain(question, text_splitter, vectorstore, retriever):
    retrieved_docs = retriever.invoke(question)
    formatted_content = combine_docs(retrieved_docs)
    return ollama_llm(question, formatted_content)

def ask_question(pdf_bytes, question):
    text_splitter, vectorstore, retriever = process_pdf(pdf_bytes)
    if text_splitter is None:
        return None
    result = rag_chain(question, text_splitter, vectorstore, retriever)
    return {result}

interface = gr.Interface(
    fn = ask_question,
    inputs = [
        gr.File(label = "上传PDF(可选)"),
        gr.Textbox(label = "请提问"),
    ],
    outputs = "text",
    title = "提一个跟上传的PDF有关的问题",
    description = "我们将用DeepSeek回答你的问题"
)

interface.launch()